"""
加attention的 GRU编码解码
"""
import random

import torch
from torch import nn
import torch.nn.functional as F
from pytorch_pretrained_bert import BertModel,BertTokenizer,BertConfig,BertAdam
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
import torch.optim as optim
import numpy as np


SEED = 4321
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)


def get_train_data(tokenizer):
    seq_len = 64
    example_pair_list = []
    # 数据
    input_text = "[CLS] 我今天心情很好[SEP] "
    target_text = "今天天气怎么样"

    example_pair = dict()

    tokenized_text = tokenizer.tokenize(input_text)
    # for i in target_text:
    #     tokenized_text.append('[MASK]')
    # # tokenized_text.append('[SEP]')
    for _ in range(seq_len - len(tokenized_text)):  # 不足128的全部补 【mask】
        tokenized_text.append('[PAD]')
    # tokenized_text.append('[MASK]')
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)

    # tokens_tensor = torch.tensor([indexed_tokens]).to(device)

    loss_ids = []
    loss_ids.append(tokenizer.convert_tokens_to_ids(['[CLS]'])[0])
    for i in target_text:
        loss_ids.append(tokenizer.convert_tokens_to_ids(i)[0])
    loss_ids.append(tokenizer.convert_tokens_to_ids(['[SEP]'])[0])
    for _ in range(seq_len - len(loss_ids)):  # 拼接长度为  128  后面的数据补-1
        loss_ids.append(0)
    # loss_tensors = torch.tensor([loss_ids]).to(device)

    # segment_ids = []
    segment_ids = [0]*seq_len

    example_pair['tokens'] = indexed_tokens
    example_pair['loss'] = loss_ids
    example_pair['segment_ids'] = segment_ids
    example_pair_list.append(example_pair)
    example_pair_list.append(example_pair)


    all_input_ids = torch.tensor([f['tokens'] for f in example_pair_list], dtype=torch.long)
    all_segment_ids = torch.tensor([f['segment_ids'] for f in example_pair_list], dtype=torch.long)
    all_labels = torch.tensor([f['loss'] for f in example_pair_list], dtype=torch.long)
    dataset = TensorDataset(all_input_ids, all_segment_ids, all_labels)

    # example_pair[tokens_tensor] = loss_tensors
    return dataset


class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
        super().__init__()

        self.embedding = nn.Embedding(input_dim, emb_dim)

        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True)

        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        # src = [src len, batch size]

        embedded = self.dropout(self.embedding(src))

        # embedded = [src len, batch size, emb dim]

        outputs, hidden = self.rnn(embedded)

        # outputs = [src len, batch size, hid dim * num directions]
        # hidden = [n layers * num directions, batch size, hid dim]

        # hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]
        # outputs are always from the last layer

        # hidden [-2, :, : ] is the last of the forwards RNN
        # hidden [-1, :, : ] is the last of the backwards RNN

        # initial decoder hidden is final hidden state of the forwards and backwards
        #  encoder RNNs fed through a linear layer
        hidden = torch.tanh(self.fc(torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1)))

        # outputs = [src len, batch size, enc hid dim * 2]
        # hidden = [batch size, dec hid dim]

        return outputs, hidden


class Attention(nn.Module):
    def __init__(self, enc_hid_dim, dec_hid_dim):
        super().__init__()

        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
        self.v = nn.Linear(dec_hid_dim, 1, bias=False)

    def forward(self, hidden, encoder_outputs):
        # hidden = [batch size, dec hid dim]
        # encoder_outputs = [src len, batch size, enc hid dim * 2]

        batch_size = encoder_outputs.shape[1]
        src_len = encoder_outputs.shape[0]

        # repeat decoder hidden state src_len times
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)

        encoder_outputs = encoder_outputs.permute(1, 0, 2)

        # hidden = [batch size, src len, dec hid dim]
        # encoder_outputs = [batch size, src len, enc hid dim * 2]

        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))

        # energy = [batch size, src len, dec hid dim]

        attention = self.v(energy).squeeze(2)

        # attention= [batch size, src len]

        return F.softmax(attention, dim=1)


class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
        super().__init__()

        self.output_dim = output_dim
        self.attention = attention

        self.embedding = nn.Embedding(output_dim, emb_dim)

        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)

        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, input, hidden, encoder_outputs):
        # input = [batch size]
        # hidden = [batch size, dec hid dim]
        # encoder_outputs = [src len, batch size, enc hid dim * 2]

        input = input.unsqueeze(0)

        # input = [1, batch size]

        embedded = self.dropout(self.embedding(input))

        # embedded = [1, batch size, emb dim]

        a = self.attention(hidden, encoder_outputs)

        # a = [batch size, src len]

        a = a.unsqueeze(1)

        # a = [batch size, 1, src len]

        encoder_outputs = encoder_outputs.permute(1, 0, 2)

        # encoder_outputs = [batch size, src len, enc hid dim * 2]

        weighted = torch.bmm(a, encoder_outputs)

        # weighted = [batch size, 1, enc hid dim * 2]

        weighted = weighted.permute(1, 0, 2)

        # weighted = [1, batch size, enc hid dim * 2]

        rnn_input = torch.cat((embedded, weighted), dim=2)

        # rnn_input = [1, batch size, (enc hid dim * 2) + emb dim]

        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))

        # output = [seq len, batch size, dec hid dim * n directions]
        # hidden = [n layers * n directions, batch size, dec hid dim]

        # seq len, n layers and n directions will always be 1 in this decoder, therefore:
        # output = [1, batch size, dec hid dim]
        # hidden = [1, batch size, dec hid dim]
        # this also means that output == hidden
        assert (output == hidden).all()

        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        weighted = weighted.squeeze(0)

        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim=1))

        # prediction = [batch size, output dim]

        return prediction, hidden.squeeze(0)


class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        # src = [src len, batch size]
        # trg = [trg len, batch size]
        # teacher_forcing_ratio is probability to use teacher forcing
        # e.g. if teacher_forcing_ratio is 0.75 we use teacher forcing 75% of the time

        batch_size = src.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim

        # tensor to store decoder outputs
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)

        # encoder_outputs is all hidden states of the input sequence, back and forwards
        # hidden is the final forward and backward hidden states, passed through a linear layer
        encoder_outputs, hidden = self.encoder(src)

        # first input to the decoder is the <sos> tokens
        input = trg[0, :]

        for t in range(1, trg_len):
            # insert input token embedding, previous hidden state and all encoder hidden states
            # receive output tensor (predictions) and new hidden state
            output, hidden = self.decoder(input, hidden, encoder_outputs)

            # place predictions in a tensor holding predictions for each token
            outputs[t] = output

            # decide if we are going to use teacher forcing or not
            teacher_force = random.random() < teacher_forcing_ratio

            # get the highest predicted token from our predictions
            top1 = output.argmax(1)

            # if teacher forcing, use actual next token as next input
            # if not, use predicted token
            input = trg[t] if teacher_force else top1

        return outputs


def init_weights(m):
    for name, param in m.named_parameters():
        if 'weight' in name:
            nn.init.normal_(param.data, mean=0, std=0.01)
        else:
            nn.init.constant_(param.data, 0)

tokenizer = BertTokenizer.from_pretrained(r'D:\work\deep_learning_model\chinese_wwm_ext_pytorch')

vocab_size = len(tokenizer.vocab)
INPUT_DIM = vocab_size
OUTPUT_DIM = vocab_size
ENC_EMB_DIM = 512
DEC_EMB_DIM = 512

ENC_HID_DIM = DEC_HID_DIM = 256
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5

attn = Attention(ENC_HID_DIM, DEC_HID_DIM)
enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)


device = torch.device('cuda')

model = Seq2Seq(enc, dec, device)

model.apply(init_weights)

model.to(device)


def train_model(model, train_dataset, clip=1):

    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=2)

    optimizer = optim.Adam(model.parameters())
    model.train()

    epoch = 50
    for _ in range(epoch):
        epoch_loss = 0
        for i, batch in enumerate(train_dataloader):
            batch = tuple(t.to(device) for t in batch)
            src = batch[0]
            trg = batch[2]
            output = model(src,trg)
            # trg = [trg len, batch size]
            # output = [trg len, batch size, output dim]

            output_dim = output.shape[-1]

            output = output[1:].view(-1, output_dim)
            trg = trg[1:].view(-1)
            # trg = [(trg len - 1) * batch size]
            # output = [(trg len - 1) * batch size, output dim]
            criterion = nn.CrossEntropyLoss(ignore_index=0)
            loss = criterion(output, trg)
            epoch_loss += loss.item()
            loss.backward()
            optimizer.step()
            model.zero_grad()
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        print(epoch_loss/(i+1))

train_dataset = get_train_data(tokenizer)
train_model(model, train_dataset)



# 9.960074424743652
# 9.945839881896973
# 9.9279203414917
# 9.903753280639648
# 9.875020027160645
# 9.838399887084961
# 9.787220001220703
# 9.71048641204834
# 9.616353034973145
# 9.472638130187988
# 9.276079177856445
# 8.966577529907227
# 8.56634521484375
# 7.963747024536133
# 7.18685245513916
# 6.073338985443115
# 4.8021063804626465
# 3.4731321334838867
# 2.6138670444488525
# 2.1472997665405273
# 1.9522149562835693
# 1.8724085092544556
# 1.8246922492980957
# 1.7949379682540894
# 1.7609933614730835
# 1.7434736490249634
# 1.7188502550125122
# 1.704430103302002
# 1.6748613119125366
# 1.649103045463562
# 1.6225777864456177
# 1.598408818244934
# 1.55029296875
# 1.5338809490203857
# 1.4860520362854004
# 1.462762475013733
# 1.423843502998352
# 1.374566674232483
# 1.3192940950393677
# 1.2714431285858154
# 1.1993308067321777
# 1.159138560295105
# 1.0942631959915161
# 1.0432343482971191
# 0.9635303020477295
# 0.8999961018562317
# 0.8483728766441345
# 0.7699136734008789
# 0.7055457234382629
# 0.6747134327888489